package pkcs5

import (
	"bytes"
	"errors"
)

func Pad(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padtext...)
}

func Unpad(data []byte, blockSize int) ([]byte, error) {
	length := len(data)
	if length == 0 || length%blockSize != 0 {
		return nil, errors.New("invalid data length")
	}
	unpadding := int(data[length-1])
	if unpadding == 0 || unpadding > blockSize {
		return nil, errors.New("invalid padding data")
	}
	padtext := data[length-unpadding:]
	for i := 0; i < unpadding; i++ {
		if padtext[i] != byte(unpadding) {
			return nil, errors.New("invalid padding data")
		}
	}
	return data[:(length - unpadding)], nil
}
